/*
 * Licensed to the Tai Ping Jin Ke
 *
 * Copyright (c) 2022 .
 * All rights reserved.
 * 项目名称：太平金科-通用能力封装层-公共组件
 * 版权说明：本软件属太平金融科技服务(上海)有限公司所有，在未获得太平金融科技服务(上海)有限公司正式授权情况下，任何企业和个人，不能获取、阅读、安装、传播本软件涉及的任何受知识产权保护的内容。
 */
package com.taipingframework.utility.thread.threadlocal;

import com.alibaba.ttl.TransmittableThreadLocal;
import org.slf4j.MDC;
import org.springframework.lang.NonNull;

import java.util.concurrent.ConcurrentHashMap;

/**
 * 实现一个线程安全的TTL
 */
public class TransmittableThreadLocalHelper {

    /**
     * 实现 TransmittableThreadLocal 的 initialValue,beforeExecute,afterExecute接口
     */
    private static final TransmittableThreadLocal<ConcurrentHashMap<String, String>> TTL_THREAD_LOCAL_ENV = new TransmittableThreadLocal<ConcurrentHashMap<String, String>>() {

        @Override
        protected ConcurrentHashMap<String, String> initialValue() {
            return new ConcurrentHashMap<>();
        }

        // 解决普通线程池中使用TTL，造成数据污染的问题
        @Override
        protected ConcurrentHashMap<String, String> childValue(ConcurrentHashMap<String, String> parentValue) {
            return initialValue();
        }

        // 父子线程使用的是拷贝对象。而非简单对象的引用。
        @Override
        public ConcurrentHashMap<String, String> copy(ConcurrentHashMap<String, String> parentValue) {
            return new ConcurrentHashMap<>(parentValue);
        }

        /**
         * 在多线程数据传递的时候，将数据复制一份给MDC
         */
        @Override
        protected void beforeExecute() {
            final ConcurrentHashMap<String, String> mdc = get();
            mdc.forEach(MDC::put);
        }

        @Override
        protected void afterExecute() {
            MDC.clear();
        }

    };

    /**
     * 设置 本地线程中 全局环境变量
     */
    public static synchronized void setEnvironment(@NonNull ConcurrentHashMap<String, String> environment) {
        TTL_THREAD_LOCAL_ENV.set(environment);
    }

    /**
     * 获取 本地线程中 全局环境变量
     */
    @NonNull
    public static ConcurrentHashMap<String, String> getEnvironment() {
        return TTL_THREAD_LOCAL_ENV.get();
    }

    /**
     * 移除 本地线程中 环境变量
     */
    public static void remove() {
        TTL_THREAD_LOCAL_ENV.remove();
    }

}
